from langgraph.graph import StateGraph, START, END
from typing_extensions import TypedDict
from langchain_ollama import ChatOllama
from langchain_core.prompts import PromptTemplate
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain.agents import Tool
from langchain.agents import initialize_agent, AgentType
import requests
import os

os.environ["TAVILY_API_KEY"] = "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"


def get_weather(location: str) -> str:
    """根据天气API 获取天气信息"""
    api_url = f"https://api.weatherapi.com/v1/current.json?key=264efe8419d84d64ab371844251707&q={location}"
    response = requests.get(api_url)
    data = response.json()
    # weather = data["current"]["condition"]["text"]
    return f"{location} 天气信息：{data}"


def get_web(query: str):
    search = TavilySearchResults()
    result = search.run(query)
    return result


my_prompt = PromptTemplate.from_template("""
尽可能回答用户的问题，你可以使用下面的工具:
"WeatherFetcher":"获得指定城市的天气JSON数据信息"
"CalcFetcher":"接收两个参数，用于计算两个数字的差"
"WebSearch":"通过Web搜索获取指定问题的信息"

严格遵守并使用下面的格式:
Question: 用户提出的问题
Thought: 思考应该如何解决问题
Action: 制定接下来要执行的动作
Action Input：接下来要执行的动作
Observation: 观察执行动作返回的结果是否能解决问题
...（Thought/Action/Action Input/Observation的过程有可能会重复多次）
Thought: 我现在可以得到最终答案了
Final Answer: 给用户输出最终答案

现在让我们开始吧!
Question: {input}
""")

weather_tool = Tool(
    name="WeatherFetcher",
    func=get_weather,
    description="用于获取指定城市的天气信息"
)

web_tool = Tool(
    name="WebSearch",
    func=get_web,
    description="通过Web搜索获取指定问题的信息"
)

llm = ChatOllama(base_url="http://10.0.2.114:11434", model="qwen3:8b", temperature=1.0)

agent = initialize_agent(
    tools=[weather_tool, web_tool],
    llm=llm,
    agent=AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True,
    handle_parsing_errors=True,
    agent_kwargs={"prompt": my_prompt},
)


class MyState(TypedDict):
    input: str
    output: str


def getChat(state: MyState):
    result = agent.run(state['input'])
    state["output"] = result
    return state


def getHtml(state: MyState):
    state['output'].replace("think", "shinn")
    return state


# 创建一个状态图
builder = StateGraph(MyState)
# 创建节点
builder.add_node('chat', getChat)
builder.add_node('html', getHtml)
# 创建边
builder.add_edge(START, 'chat')
builder.add_edge("chat", 'html')
builder.add_edge("chat", END)
# 编译图
graph = builder.compile()

# 执行图
result_state = graph.invoke({"input": "JS如何生成随机数", "output": ""})

# 打印结果
print(result_state)

